package cn.tworice.cryption.util;

import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.params.ECKeyGenerationParameters;
import org.bouncycastle.crypto.params.ECPrivateKeyParameters;
import org.bouncycastle.crypto.params.ECPublicKeyParameters;
import org.bouncycastle.crypto.params.ParametersWithRandom;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.SecretKeySpec;
import java.security.SecureRandom;
import java.util.Base64;

public class SMHybridCryptoUtil {

    // 生成随机 SM4 密钥（128 位）
    public static byte[] generateSM4Key() throws Exception {
        KeyGenerator kg = KeyGenerator.getInstance("SM4", "BC");
        kg.init(128);
        return kg.generateKey().getEncoded();
    }

    // 使用 SM4 加密数据
    public static byte[] sm4Encrypt(byte[] key, byte[] data) throws Exception {
        Cipher cipher = Cipher.getInstance("SM4/ECB/PKCS5Padding", "BC");
        SecretKeySpec keySpec = new SecretKeySpec(key, "SM4");
        cipher.init(Cipher.ENCRYPT_MODE, keySpec);
        return cipher.doFinal(data);
    }

    // 使用 SM4 解密数据
    public static byte[] sm4Decrypt(byte[] key, byte[] cipherText) throws Exception {
        Cipher cipher = Cipher.getInstance("SM4/ECB/PKCS5Padding", "BC");
        SecretKeySpec keySpec = new SecretKeySpec(key, "SM4");
        cipher.init(Cipher.DECRYPT_MODE, keySpec);
        return cipher.doFinal(cipherText);
    }

    // 使用 SM2 加密 SM4 密钥
    public static byte[] sm2EncryptSM4Key(byte[] sm4Key, ECPublicKeyParameters publicKey) throws Exception {
        SM2Engine sm2 = new SM2Engine();
        sm2.init(true, new ParametersWithRandom(publicKey, new SecureRandom()));
        return sm2.processBlock(sm4Key, 0, sm4Key.length);
    }

    // 使用 SM2 解密 SM4 密钥
    public static byte[] sm2DecryptSM4Key(byte[] encryptedKey, ECPrivateKeyParameters privateKey) throws Exception {
        SM2Engine sm2 = new SM2Engine();
        sm2.init(false, privateKey);
        return sm2.processBlock(encryptedKey, 0, encryptedKey.length);
    }

    // 混合加密：返回 Base64 编码后的 sm2EncryptedKey 和 sm4CipherText
    public static HybridEncryptedData hybridEncrypt(String plainText, ECPublicKeyParameters sm2PublicKey) throws Exception {
        byte[] sm4Key = generateSM4Key();
        byte[] cipherText = sm4Encrypt(sm4Key, plainText.getBytes("UTF-8"));
        byte[] encryptedSm4Key = sm2EncryptSM4Key(sm4Key, sm2PublicKey);

        return new HybridEncryptedData(
                Base64.getEncoder().encodeToString(encryptedSm4Key),
                Base64.getEncoder().encodeToString(cipherText)
        );
    }

    // 混合解密
    public static String hybridDecrypt(HybridEncryptedData data, ECPrivateKeyParameters sm2PrivateKey) throws Exception {
        byte[] encryptedSm4Key = Base64.getDecoder().decode(data.getEncryptedSm4Key());
        byte[] sm4CipherText = Base64.getDecoder().decode(data.getCipherText());

        byte[] sm4Key = sm2DecryptSM4Key(encryptedSm4Key, sm2PrivateKey);
        byte[] plainBytes = sm4Decrypt(sm4Key, sm4CipherText);

        return new String(plainBytes, "UTF-8");
    }

    // 封装密文结构体
    public static class HybridEncryptedData {
        private String encryptedSm4Key;
        private String cipherText;

        public HybridEncryptedData(String encryptedSm4Key, String cipherText) {
            this.encryptedSm4Key = encryptedSm4Key;
            this.cipherText = cipherText;
        }

        public String getEncryptedSm4Key() {
            return encryptedSm4Key;
        }

        public String getCipherText() {
            return cipherText;
        }
    }
}
